#include <omp.h>
#include <vector>
#include <cstring>

#include "no-atomic.hh"
#include "../utils/impl.hh"
#include "../utils/debug.hh"
#include "../utils/timer.hh"

namespace icesp
{
    namespace kernal
    {
        int no_atomic::nthreads;
        int no_atomic::nrelax;
        int const* no_atomic::colors;
        heaps* no_atomic::_heaps;
        no_atomic::twod_cqueue_type* no_atomic::queues;
        volatile int* no_atomic::versions;
        std::string const no_atomic::name{"no-atomic"};

        void no_atomic::relax(int pid, graph::adj_matrix& matrix)
        {
            auto const& edges      = matrix.edges;
            auto const& edge_index = matrix.edge_index;
            auto      & result     = matrix.result;

            impl::item data;
            auto& src_node   = data.key;
            auto& src_result = data.value;

            auto left = nrelax;
            while (left > 0 && !_heaps->empty(pid)) {
                _heaps->remove_top(pid, &data);
                auto push_ret = true;
                for (int i = edge_index[src_node]; i < edge_index[src_node + 1]; i++) {
                    auto dst_node = edges[i].to;
                    auto weight = edges[i].weight;
                    auto new_dst_result = src_result + weight;

                    /* case colors[dst_node] == pid, we get right result[dst_node] */
                    /* case colors[dst_node] != pid, current thread may have a cache value, but the cache value must be bigger than the actural one. */

                    if (colors[dst_node] == pid && result[dst_node] > new_dst_result) {
                        result[dst_node] = new_dst_result;
                        _heaps->add_or_update(pid, dst_node, new_dst_result);
                    } else if (colors[dst_node] != pid) {
                        push_ret = false;
                        small_concurrent_queue<>* dst_queue = queues->get(colors[dst_node], pid);
                        for (int tmp = 0; tmp < utils::push_retry_times; tmp++) {
                            push_ret = dst_queue->push(impl::item{dst_node, new_dst_result});
                            if (push_ret) break;
                        }
                        if (!push_ret) break;
                    }
                }

                if (!push_ret) {
                    /* push fails, go to adjust and later retry relax */
                    _heaps->add(pid, src_node, src_result);
                    break;
                } else
                    left--;
            }
        }

        void no_atomic::adjust(int pid, graph::adj_matrix& mat)
        {
            auto & result = mat.result;
            impl::item data;
            auto& src_node   = data.key;
            auto& src_result = data.value;
            for (auto i = 0; i < nthreads; i++) {
                while (!queues->get(pid, i)->empty()) {
                    data = queues->get(pid, i)->pop();
                    if (result[src_node] > src_result) {
                        result[src_node] = src_result;
                        _heaps->add_or_update(pid, src_node, src_result);
                    }
                }
            }
        }

        bool no_atomic::update_queue_empty(int pid)
        {
            for (auto i = 0; i < nthreads; i++)
                if (!queues->get(pid, i)->empty() || !queues->get(i, pid)->empty())
                    return false;
            return true;
        }

        bool no_atomic::heap_queue_empty(int pid)
        {
            return _heaps->empty(pid) && update_queue_empty(pid);
        }

        void no_atomic::worker(graph::adj_matrix& mat)
        {
            auto const pid = omp_get_thread_num();
            auto& my_version = ((int*) versions)[pid];
            int local_versions[nthreads];
            auto more = true;
            while (more) {
                if (my_version >= 0) {
                    relax(pid, mat);
                    adjust(pid, mat);
                    if (heap_queue_empty(pid))
                        my_version = ~((my_version + 1) & 0x80000000);
                } else {
                    if (!heap_queue_empty(pid)) {
                        my_version = ~my_version;
                        continue;
                    }
                    auto some_state_active = false;
                    for (auto i = 0; i < nthreads && !some_state_active; i++) {
                        if (versions[i] >= 0) some_state_active = true;
                        local_versions[i] = versions[i];
                    }
                    if (some_state_active) continue;
                    auto is_finished = true;
                    for (auto i = 0; i < nthreads && !some_state_active && is_finished; i++) {
                        if (versions[i] >= 0) some_state_active = true;
                        if (local_versions[i] != versions[i]) is_finished = false;
                    }
                    if (some_state_active) continue;
                    if (is_finished) more = false;
                    break;
                }
            }
        }

        double no_atomic::compute(graph::adj_matrix& mat, painter& p, int source, int target)
        {
            PRINT3("", name << ": ", "shortest path compute started.\n");

            omp_set_num_threads(nthreads);

            auto& n = mat.n;
            colors = p.get_colors();

            auto& result = mat.result;
            mat.reset_result();

            // create heaps and queues for each thread
            // FIXME impl colors
            _heaps = new struct heaps(n, nthreads, p.get_node_number_count());
            queues = new twod_cqueue_type(nthreads);

            versions = new int[sizeof(int) * nthreads];
            std::memset((void*)versions, 0, sizeof(int) * nthreads);

            timer running_timer;
            running_timer.start();

            result[source] = 0;
            _heaps->add(colors[source], source, 0);

#pragma omp parallel
            {
                worker(mat);
            }

            running_timer.stop();

            delete[] versions;
            delete queues;
            delete _heaps;

            // DEBUG("result[target = " << target << "] = " << mat.result[target]);

            return running_timer.elapsed_seconds();
        }

    }
}

